import torch
import torch.optim as optim
from torch import nn
from tqdm import tqdm  # 导入tqdm用于进度显示

from Dataloader import MedicalDataset
from TransUnet import TransUNet


def train(model, dataloader, criterion, optimizer, device, num_epochs=31):
    model.to(device)  # 将模型移动到指定设备
    model.train()  # 设置模型为训练模式

    # 使用tqdm包装epoch循环，设置position=0使其显示在顶部
    epoch_pbar = tqdm(range(num_epochs), desc='Training Epochs', position=0)
    for epoch in epoch_pbar:
        epoch_loss = 0.0  # 累计每个epoch的损失

        # 使用tqdm包装dataloader，position=1使其显示在epoch进度条下方
        batch_pbar = tqdm(enumerate(dataloader),
                          total=len(dataloader),
                          desc=f'Epoch {epoch + 1}/{num_epochs}',
                          position=1,
                          leave=False)  # leave=False确保每个epoch结束后清除该进度条

        for batch_idx, (images, masks) in batch_pbar:
            images = images.float().to(device)  # 将图像移动到设备并转换为浮点数
            masks = masks.float().to(device)  # 将掩码移动到设备并转换为浮点数

            optimizer.zero_grad()  # 清零梯度
            outputs = model(images)  # 前向传播

            masks = masks.unsqueeze(dim=1).float().to(device)
            loss = criterion(outputs, masks)  # 计算损失

            loss.backward()  # 反向传播
            optimizer.step()  # 更新参数

            epoch_loss += loss.item()  # 累加损失

            # 更新进度条显示当前batch的loss
            batch_pbar.set_postfix({'batch_loss': f'{loss.item():.4f}'}, refresh=True)

        # 关闭batch进度条
        batch_pbar.close()

        # 计算并显示每个epoch的平均损失
        average_loss = epoch_loss / len(dataloader)
        epoch_pbar.set_postfix({'avg_loss': f'{average_loss:.4f}'})

        # 每隔3个周期保存一次模型
        if (epoch + 1) % 3 == 0:
            save_path = f'../pt_file/Unet_Upsample-{epoch + 1}.pt'
            torch.save(model.state_dict(), save_path)
            tqdm.write(f'Model saved to {save_path}')


if __name__ == '__main__':
    # 创建训练数据集
    train_dataset = MedicalDataset(
        root_dir='/Volumes/For_Mac/dateset/Pulmonary_X_ray_and_masks',
        is_train=True,
        image_size=256
    )

    # 创建测试数据集
    test_dataset = MedicalDataset(
        root_dir='/Volumes/For_Mac/dateset/Pulmonary_X_ray_and_masks',
        is_train=False,
        image_size=256
    )

    # 创建数据加载器
    from torch.utils.data import DataLoader

    train_loader = DataLoader(
        train_dataset,
        batch_size=4,
        shuffle=True,
        num_workers=4
    )

    test_loader = DataLoader(
        test_dataset,
        batch_size=1,
        shuffle=False,
        num_workers=4
    )

    # 初始化模型、损失函数和优化器
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = TransUNet(num_classes=1).to(device)

    criterion = nn.BCELoss()  # 是否使用BCELoss要根据输出是否使用了sigmoid函数来判断，如使用了则使用BCELoss
    optimizer = optim.Adam(model.parameters(), lr=0.001)

    # 开始训练
    train(model, train_loader, criterion, optimizer, device)
